package com.gdut.laiagent.rag;

import jakarta.annotation.Resource;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.core.JdbcTemplate;

import java.util.List;

import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgDistanceType.COSINE_DISTANCE;
import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgIndexType.HNSW;

//@Configuration
public class PgVectorVectorStoreConfig {

    // 引入文档加载器，把文档加载到向量存储中
    @Resource
    LoveAppDocumentLoader loveAppDocumentLoader;

//    @Bean
    VectorStore pgVectorVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel dashscopeEmbeddingModel){
        PgVectorStore vectorStore = PgVectorStore.builder(jdbcTemplate, dashscopeEmbeddingModel)
                .dimensions(1536)                    // Optional: defaults to model dimensions or 1536
                .distanceType(COSINE_DISTANCE)       // Optional: defaults to COSINE_DISTANCE,表示使用余弦距离
                .indexType(HNSW)                     // Optional: defaults to HNSW,高性能的近似最近邻（ANN）搜索算法索引，用于加速相似性搜索的索引类型。
                .initializeSchema(true)              // Optional: defaults to false,决定是否自动初始化数据库模式（Schema）,自动创建所需的表（vector_table_name 指定的表）和索引
                .schemaName("public")                // Optional: defaults to "public",设置存储向量表所在的数据库模式（Schema）,表将被创建在 public 模式下。
                .vectorTableName("vector_store")     // Optional: defaults to "vector_store",设置存储向量和元数据的数据库表名。
                .maxDocumentBatchSize(10000)         // Optional: defaults to 10000,设置在批量处理文档时（例如初始化添加大量文档），每批处理的文档数量上限。
                .build();
        // 加载文档
        List<Document> documents = loveAppDocumentLoader.loadMarkdowns();
        vectorStore.add(documents);
        return vectorStore;
    }
}
